package netcom;

import java.io.*;
import java.net.*;
import java.security.InvalidKeyException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.PrivateKey;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;

public class Protocol
{
    private boolean encrypt;
    private KeyPair keyPair;
    private PublicKey encryptKey;
    private PrivateKey decryptKey;
    private Socket socket;
    private ObjectOutputStream out;
    private ObjectInputStream in;
    private Package inData;
    private Package outData;
    private static final String ALGORITHM = "RSA";
    private static final String ALG_PREF = "/ECB/NoPadding";
    private static final int KEYSIZE = 1024;
    private static final String ENCRYPT = "ENCRYPT";
    private static final String ENDCODE = "ENDCODE";
    private static final String YES = "YES";
    private static final String NO = "NO";
    private Cipher cipher;
    private String endCode;
    private boolean negotiation_done;
    private boolean negotiating;
    private boolean debug;

    public Protocol()
    {
        keyPair = null;
        encrypt = false;
        socket = null;
        outData = new Package();
        inData = new Package();
        cipher = null;
        negotiation_done = false;
        negotiating = true;
        debug = false;
    }
    public Protocol(KeyPair keyPair)
    {
        this.keyPair = keyPair;
        this.encrypt = false;
        socket = null;
        outData = new Package();
        inData = new Package();
        cipher = null;
        negotiation_done = false;
        negotiating = true;
        debug = false;
    }
    public void setKeys(KeyPair keyPair)
    {
        this.keyPair = keyPair;
    }
    public void setSocket(Socket socket)
    {
        this.socket = socket;
        try
        {
                out = new ObjectOutputStream(this.socket.getOutputStream());    // output stream must go as first, otherwise it stucks Java (sic!)
                in = new ObjectInputStream(this.socket.getInputStream());
        }
        catch (IOException e)
        {
                System.out.println("Could not get input or output stream");
                return;
        }
    }
    public String getEndCode()
    {
        return endCode;
    }
    public void negotiateWithServer(Socket sock)
    {
        negotiating = true;
        System.out.println("Starting to negotiate");
        setSocket(sock);
        if(socket == null)
        {
            System.out.println("Internal error: server socket not opened");
            return;
        }
        // Try to negotiate encryption mode
        // send query for encryption
        sendMsg(ENCRYPT);
        String answer = getMsg();
        System.out.println("Encrypting data: " + answer);
        // if positive generate keys and exchange them
        if(answer.compareTo(YES) == 0)
        {
            encrypt = true;
            try
            {
                cipher = Cipher.getInstance(ALGORITHM + ALG_PREF);
            }
            catch (NoSuchAlgorithmException ex)
            {
                System.out.println("Wrong encryption algorithm. Shutting down(protocol)...");
                System.exit(-1);
            }
            catch (NoSuchPaddingException ex)
            {
                System.out.println("No such padding. Shutting down(protocol)...");
                System.exit(-1);
            }
            try
            {
                try
                {
                    System.out.print("Generating keys: ");
                    KeyPairGenerator keypairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
                    keypairGenerator.initialize(KEYSIZE);
                    keyPair = keypairGenerator.generateKeyPair();
                    System.out.println("DONE");
                }
                catch(NoSuchAlgorithmException e)
                {
                    System.out.println("Wrong encryption algorithm. Shutting down(protocol)...");
                    System.exit(-1);
                }
                // key exchange
                // get server public key
                System.out.print("Key exchange: ");
                try
                {
                    this.encryptKey = (PublicKey)in.readObject();
                }
                catch (IOException ex)
                {
                    System.out.println("Could not get Public Key");
                }
                this.decryptKey = keyPair.getPrivate();
                // send this public key
                try
                {
                    out.writeObject(keyPair.getPublic());
                }
                catch(IOException ex)
                {
                    System.out.println("Could not send Public Key");
                }
                System.out.println("DONE");
            }
            catch(ClassNotFoundException ex)
            {
                System.out.println("Class not found");
                return;
            }
        }

        System.out.println("Getting end code");
        // Try to get end code
        sendMsg(ENDCODE);
        endCode = getMsg();
        System.out.println("Recieved end code: " + endCode);
        System.out.println("Negotiation: DONE");
        negotiation_done = true;
        negotiating = false;
    }

    public void negotiateWithClient(Socket sock)
    {
        negotiating = true;
        System.out.println("Starting to negotiate");
        setSocket(sock);
        if(socket == null)
        {
            System.out.println("Internal error: client socket not opened");
            return;
        }
        // check if server has set keys (encryption on)
        if(keyPair != null)
        {
            decryptKey = keyPair.getPrivate();
            try
            {
                cipher = Cipher.getInstance(ALGORITHM + ALG_PREF);
            }
            catch (NoSuchAlgorithmException ex)
            {
                System.out.println("Wrong encryption algorithm. Shutting down(protocol)...");
                System.exit(-1);
            }
            catch (NoSuchPaddingException ex)
            {
                System.out.println("No such padding. Shutting down(protocol)...");
                System.exit(-1);
            }
        }
        if(getMsg().compareTo(ENCRYPT) == 0)
        {
            if(keyPair != null)
            {
                try
                {
                    sendMsg(YES);
                    System.out.print("Key exchange: ");
                    out.writeObject(keyPair.getPublic());
                    try
                    {
                        encryptKey = (PublicKey)in.readObject();
                    }
                    catch(ClassNotFoundException ex)
                    {
                        System.out.println("Class not found");
                    }
                    decryptKey = keyPair.getPrivate();
                    System.out.println("DONE");
                }
                catch(IOException ex)
                {
                    System.out.println("Could not send PublicKey");
                }
                encrypt = true;
            }
            else
                sendMsg(NO);
        }
        if(getMsg().compareTo(ENDCODE) == 0)
        {
            endCode = ENDCODE.toString();
            endCode = String.valueOf(endCode.hashCode());
            System.out.println("Sending end code: " + endCode);
            sendMsg(endCode);
        }
        System.out.println("Negotiation: DONE");
        negotiation_done = true;
        negotiating = false;
    }

    public void sendMsg(String in_)
    {
        if(negotiation_done == false && negotiating == false)
        {
            System.out.println("Try to negotiate the connection first!");
            return;
        }
        byte[] input = null;
        if(out == null)
        {
            System.out.println("No connection. Did you negotiate with the server ?");
            return;
        }
        System.out.println("Sending: " + in_);
        if(encrypt)
        {
            // code spaces
            input = (in_.replace(' ', '¤')).getBytes();
            try
            {
                cipher.init(Cipher.ENCRYPT_MODE, encryptKey);
            }
            catch (InvalidKeyException ex)
            {
                System.out.println("Invalid key. No encryption.");
                encrypt = false;
            }
            if(encrypt)
            {
                int textBytesChunkLen = 100;
                int encryptedChunkNum = (input.length - 1) / textBytesChunkLen + 1;
                // RSA returns 128 bytes as output for 100 text bytes
                int encryptedBytesChunkLen = 128;
                int encryptedBytesLen = encryptedChunkNum * encryptedBytesChunkLen;
                // define the output array.
                byte[] encryptedBytes = new byte[encryptedBytesLen];
                int textBytesChunkIndex = 0;
                int encryptedBytesChunkIndex = 0;
                try
                {
                    for (int i = 0; i < encryptedChunkNum; i++)
                    {
                        try
                        {
                            if (i < encryptedChunkNum - 1)
                            {
                                encryptedBytesChunkIndex = encryptedBytesChunkIndex
                                        + cipher.doFinal(input, textBytesChunkIndex, textBytesChunkLen, encryptedBytes, encryptedBytesChunkIndex);
                                textBytesChunkIndex = textBytesChunkIndex + textBytesChunkLen;
                            }
                            else
                            {
                                cipher.doFinal(input, textBytesChunkIndex,input.length - textBytesChunkIndex,encryptedBytes, encryptedBytesChunkIndex);
                            }
                        }
                        catch(ShortBufferException ex)
                        {
                                System.out.println("Short buffer error");
                        }
                    }
                }
                catch(IllegalBlockSizeException ex)
                {
                    System.out.println("Illegal block size.");
                }
                catch(BadPaddingException ex)
                {
                    System.out.println("Bad padding.");
                }
                input = encryptedBytes;
            }
            if(debug)
            {
                String tmp = new String(input);
                System.out.println("Sending(encrypted): " + tmp);
            }
        }
        else
        {
            input = in_.getBytes();
        }
        outData.setData(input);
        try
        {
            out.reset();
            out.writeObject(outData);
            out.flush();
        }
        catch (IOException ex)
        {
            ex.printStackTrace(System.out);
            System.out.println("Could not send data");
        }
    }

    public String getMsg()
    {
        if(negotiation_done == false && negotiating == false)
        {
            System.out.println("Try to negotiate the connection first!");
            return null;
        }
        String result = null;
        if(in == null)
        {
            System.out.println("No connection. Did you negotiate with the server ?");
            return null;
        }
        try
        {
            try
            {
                inData = (Package)in.readObject();
                result = new String(inData.getData());
                if(debug && encrypt)
                    System.out.println("Reading: " + result);
            }
            catch (IOException ex)
            {
                System.out.println("Could not read data");
                return null;
            }
        }
        catch (ClassNotFoundException ex)
        {
            System.out.println("Class not found (protocol get): " + ex);
            return null;
        }
        if(encrypt)
        {
            try
            {
                cipher.init(Cipher.DECRYPT_MODE, decryptKey);
            }
            catch (InvalidKeyException ex)
            {
                System.out.println("Invalid key. No encryption.");
                encrypt = false;
            }
            if(encrypt)
            {
                result = "";
                try
                {
                    byte[] encryptedBytes = inData.getData();
                    int encryptedByteChunkLen = 128;
                    int encryptedChunkNum = encryptedBytes.length / encryptedByteChunkLen;
                    byte[] decryptedBytesChunk = new byte[encryptedByteChunkLen];
                    int decryptedIndex = 0;
                    int encryptedIndex = 0;
                    for (int i = 0; i < encryptedChunkNum; i++)
                    {
                        try
                        {
                            if (i < encryptedChunkNum - 1)
                            {
                                cipher.doFinal(encryptedBytes, encryptedIndex, encryptedByteChunkLen, decryptedBytesChunk);
                                encryptedIndex = encryptedIndex + encryptedByteChunkLen;
                                //String tmp = new String(decryptedBytesChunk,28,100);
                                String tmp = new String(decryptedBytesChunk);
                                result = result + tmp.trim();
                                if(debug)
                                    System.out.println(tmp);
                            }
                            else
                            {
                                decryptedIndex = cipher.doFinal(encryptedBytes, encryptedIndex, encryptedBytes.length - encryptedIndex, decryptedBytesChunk);
                                String tmp = new String(decryptedBytesChunk);
                                if(debug)
                                    System.out.println(tmp);
                                result = result + tmp.trim();
                            }
                        }
                        catch(ShortBufferException ex)
                        {
                            System.out.println("Short buffer error");
                        }
                    }
                    // decode spaces
                    result = result.replace('¤', ' ');
                }
                catch(IllegalBlockSizeException ex)
                {
                    System.out.println("Illegal block size.");
                }
                catch(BadPaddingException ex)
                {
                    System.out.println("Bad padding.");
                }
            }
            System.out.println("Reading(decrypted): " + result);
        }
        return result;
    }
    public void close()
    {
        sendMsg(endCode);
        try
        {
            in.close();
            out.close();
        }
        catch (IOException ex)
        {
            System.out.println("Could not close input or output stream");
        }
        try
        {
            socket.close();
        }
        catch (IOException ex)
        {
            System.out.println("Could not properly close socket");
        }
    }
}
